#include "tcpclient.h"
#include <cerrno>
#include <cassert>
#include <cstring>
#include <arpa/inet.h>
#include <fcntl.h>
#include <unistd.h>

static bool set_non_block(int sockfd) {
    int opts;
    if ((opts = fcntl(sockfd, F_GETFL)) < 0) {
        return false;
    }
    opts = opts | O_NONBLOCK;
    if (fcntl(sockfd, F_SETFL, opts) < 0) {
        return false;
    }
    return true;
}

static int select_read(int sockfd, int msec) {
    struct timeval tv;
    tv.tv_sec = msec / 1000;
    tv.tv_usec = msec % 1000 * 1000;

    fd_set rfs;
    FD_ZERO(&rfs);
    FD_SET(sockfd, &rfs);

    int err = select(sockfd + 1, &rfs, NULL, NULL, &tv);
    if (err < 0) {
        return TCP_CLIENT_E_SELECT;
    }

    //select timeout
    if (err == 0) {
        return TCP_CLIENT_E_TIMEOUT;
    }

    //socket no in fd_set
    if (!FD_ISSET(sockfd, &rfs)) {
        return TCP_CLIENT_E_FD_ISSET;
    }

    return TCP_CLIENT_E_OK;
}

static int select_write(int sockfd, int msec) {
    struct timeval tv;
    tv.tv_sec = msec / 1000;
    tv.tv_usec = msec % 1000 * 1000;

    fd_set wfs;
    FD_ZERO(&wfs);
    FD_SET(sockfd, &wfs);

    int err = select(sockfd + 1, NULL, &wfs, NULL, &tv);
    if (err < 0) {
        return TCP_CLIENT_E_SELECT;
    }

    //select timeout
    if (err == 0) {
        return TCP_CLIENT_E_TIMEOUT;
    }

    //socket no in fd_set
    if (!FD_ISSET(sockfd, &wfs)) {
        return TCP_CLIENT_E_FD_ISSET;
    }

    return TCP_CLIENT_E_OK;
}

static int select_both(int sockfd, int msec) {
    struct timeval tv;
    tv.tv_sec = msec / 1000;
    tv.tv_usec = msec % 1000 * 1000;

    fd_set rfs, wfs;
    FD_ZERO(&rfs);
    FD_ZERO(&wfs);
    FD_SET(sockfd, &rfs);
    FD_SET(sockfd, &wfs);

    int err = select(sockfd + 1, &rfs, &wfs, NULL, &tv);
    if (err < 0) {
        return TCP_CLIENT_E_SELECT;
    }

    //select timeout
    if (err == 0) {
        return TCP_CLIENT_E_TIMEOUT;
    }

    //socket no in fd_set
    if (!FD_ISSET(sockfd, &rfs) && !FD_ISSET(sockfd, &wfs)) {
        return TCP_CLIENT_E_FD_ISSET;
    }

    return TCP_CLIENT_E_OK;
}

typedef int(*tcp_trans_func)(int sockfd, char* buf, size_t size);
typedef int(*tcp_select_func)(int sockfd, int msec);

static int tcp_read(int sockfd, char* buf, size_t size) {
    return ::read(sockfd, buf, size);
}

static int tcp_write(int sockfd, char* buf, size_t size) {
    return ::write(sockfd, buf, size);
}

static int tcp_trans(int sockfd, char* buf, size_t size, int msec, tcp_trans_func trans_func, tcp_select_func select_func) {
    assert(buf);
    assert(msec > 0);

    int total = 0;
    while (total < size) {
        int n = trans_func(sockfd, buf + total, size - total);
        //收发数据成功
        if (n > 0) {
            total += n;
            continue;
        }
        //连接关闭
        if (n == 0) {
            return TCP_CLIENT_E_EOF;
        }
        //错误码判断
        switch (errno) {
        case EAGAIN: //重试
            {
                int err = select_func(sockfd, msec);
                if (err != TCP_CLIENT_E_OK) {
                    return err;
                }
            }
            break;
        case EINTR: //信号中断
            break;
        default: //错误
            return TCP_CLIENT_E_TRANS;
        }
    }
    return TCP_CLIENT_E_OK;
}

TcpClient::TcpClient() 
: sock_fd_(-1) {
}

TcpClient::~TcpClient() {
    disconnect();
}

int TcpClient::connect(const char* ip, int port, int msec) {
    assert(ip);
    assert(port > 0);

    bzero(&sock_addr_, sizeof(sock_addr_));
    sock_addr_.sin_family = AF_INET;
    sock_addr_.sin_port = htons(port);
    inet_aton(ip, &sock_addr_.sin_addr);

    return reconnect(msec);
}

int TcpClient::reconnect(int msec) {
    disconnect();

    sock_fd_ = socket(AF_INET, SOCK_STREAM, 0);
    if (sock_fd_ < 0) {
        return TCP_CLIENT_E_SOCKET;
    }

    //设置为非阻塞
    if (!set_non_block(sock_fd_)) {
        return TCP_CLIENT_E_SET_NON_BLOCK;
    }

    socklen_t len = sizeof(sock_addr_);
    if (::connect(sock_fd_, (struct sockaddr*)&sock_addr_, len) == 0) {
        return TCP_CLIENT_E_OK;
    }

    if (errno != EINPROGRESS) {
        return TCP_CLIENT_E_CONNECT;
    }

    int err = select_both(sock_fd_, msec);
    if (err != TCP_CLIENT_E_OK) {
        return err;
    }

    len = sizeof(err);
    if (getsockopt(sock_fd_, SOL_SOCKET, SO_ERROR, &err, &len) < 0) {
        return TCP_CLIENT_E_GET_SOCK_OPT;
    }

    if (err != 0) {
        return TCP_CLIENT_E_SOCK_OPT;
    }

    return TCP_CLIENT_E_OK;
}

void TcpClient::disconnect() {
    if (sock_fd_ != -1) {
        shutdown(sock_fd_, 2);
        close(sock_fd_);

        sock_fd_ = -1;
    }
}

int TcpClient::recv(char* buf, int size, int msec) {
    return tcp_trans(sock_fd_, buf, size, msec, tcp_read, select_read);
}

int TcpClient::send(const char* buf, int size, int msec) {
    return tcp_trans(sock_fd_, const_cast<char*>(buf), size, msec, tcp_write, select_write);
}

const char* TcpClient::get_errstr(int err) {
    switch (err) {
    case TCP_CLIENT_E_OK:
        return "ok";
    case TCP_CLIENT_E_TIMEOUT:
        return "timeout";
    case TCP_CLIENT_E_SOCKET:
        return "socket fail";
    case TCP_CLIENT_E_SET_NON_BLOCK:
        return "set non block fail";
    case TCP_CLIENT_E_CONNECT:
        return "connect fail";
    case TCP_CLIENT_E_SELECT:
        return "select fail";
    case TCP_CLIENT_E_FD_ISSET:
        return "select not in fd_set";
    case TCP_CLIENT_E_GET_SOCK_OPT:
        return "get sock opt fail";
    case TCP_CLIENT_E_SOCK_OPT:
        return "sock opt error";
    case TCP_CLIENT_E_TRANS:
        return "recv/send fail";
    case TCP_CLIENT_E_EOF:
        return "recv/send return zero";
    default:
        return "undefined error";
    }
}
